Add Improved AutoEncoder model.

plugins/Model_IAE/AutoEncoder.py

@@ -0,0 +1,46 @@
+# Improved-AutoEncoder base classes
+
+
+encoderH5 = 'encoder.h5'
+decoderH5 = 'decoder.h5'
+inter_AH5 = 'inter_A.h5'
+inter_BH5 = 'inter_B.h5'
+inter_bothH5 = 'inter_both.h5'
+
+
+class AutoEncoder:
+    def __init__(self, model_dir):
+        self.model_dir = model_dir
+
+        self.encoder = self.Encoder()
+        self.decoder = self.Decoder()
+        self.inter_A = self.Intermidiate()
+        self.inter_B = self.Intermidiate()
+        self.inter_both = self.Intermidiate()
+
+        self.initModel()
+
+    def load(self, swapped):
+        (face_A,face_B) = (inter_AH5, inter_BH5) if not swapped else (inter_AH5, inter_BH5)
+
+        try:
+            self.encoder.load_weights(str(self.model_dir / encoderH5))
+            self.decoder.load_weights(str(self.model_dir / decoderH5))
+            self.inter_both.load_weights(str(self.model_dir / inter_bothH5))
+            self.inter_A.load_weights(str(self.model_dir / face_A))
+            self.inter_B.load_weights(str(self.model_dir / face_B))
+            print('loaded model weights')
+            return True
+        except Exception as e:
+            print('Failed loading existing training data.')
+            print(e)
+            return False
+
+    def save_weights(self):
+        self.encoder.save_weights(str(self.model_dir / encoderH5))
+        self.decoder.save_weights(str(self.model_dir / decoderH5))
+        self.inter_both.save_weights(str(self.model_dir / inter_bothH5))
+        self.inter_A.save_weights(str(self.model_dir / inter_AH5))
+        self.inter_B.save_weights(str(self.model_dir / inter_BH5))
+        print('saved model weights')
+

plugins/Model_IAE/Model.py

@@ -0,0 +1,72 @@
+# Improved autoencoder for faceswap.
+
+from keras.models import Model as KerasModel
+from keras.layers import Input, Dense, Flatten, Reshape, Concatenate
+from keras.layers.advanced_activations import LeakyReLU
+from keras.layers.convolutional import Conv2D
+from keras.optimizers import Adam
+
+from .AutoEncoder import AutoEncoder
+from lib.PixelShuffler import PixelShuffler
+
+IMAGE_SHAPE = (64, 64, 3)
+ENCODER_DIM = 1024
+
+
+class Model(AutoEncoder):
+    def initModel(self):
+        optimizer = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
+        x = Input(shape=IMAGE_SHAPE)
+
+        self.autoencoder_A = KerasModel(x, self.decoder(Concatenate()([self.inter_A(self.encoder(x)), self.inter_both(self.encoder(x))])))
+        self.autoencoder_B = KerasModel(x, self.decoder(Concatenate()([self.inter_B(self.encoder(x)), self.inter_both(self.encoder(x))])))
+
+        self.autoencoder_A.compile(optimizer=optimizer, loss='mean_absolute_error')
+        self.autoencoder_B.compile(optimizer=optimizer, loss='mean_absolute_error')
+
+    def converter(self, swap):
+        autoencoder = self.autoencoder_B if not swap else self.autoencoder_A
+        return lambda img: autoencoder.predict(img)
+
+    def conv(self, filters):
+        def block(x):
+            x = Conv2D(filters, kernel_size=5, strides=2, padding='same')(x)
+            x = LeakyReLU(0.1)(x)
+            return x
+        return block
+
+    def upscale(self, filters):
+        def block(x):
+            x = Conv2D(filters * 4, kernel_size=3, padding='same')(x)
+            x = LeakyReLU(0.1)(x)
+            x = PixelShuffler()(x)
+            return x
+        return block
+
+    def Encoder(self):
+        input_ = Input(shape=IMAGE_SHAPE)
+        x = input_
+        x = self.conv(128)(x)
+        x = self.conv(256)(x)
+        x = self.conv(512)(x)
+        x = self.conv(1024)(x)
+        x = Flatten()(x)
+        return KerasModel(input_, x)
+
+    def Intermidiate(self):
+        input_ = Input(shape=(None, 4 * 4 * 1024))
+        x = input_
+        x = Dense(ENCODER_DIM)(x)
+        x = Dense(4 * 4 * int(ENCODER_DIM/2))(x)
+        x = Reshape((4, 4, int(ENCODER_DIM/2)))(x)
+        return KerasModel(input_, x)
+
+    def Decoder(self):
+        input_ = Input(shape=(4, 4, ENCODER_DIM))
+        x = input_
+        x = self.upscale(512)(x)
+        x = self.upscale(256)(x)
+        x = self.upscale(128)(x)
+        x = self.upscale(64)(x)
+        x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
+        return KerasModel(input_, x)

plugins/Model_IAE/Trainer.py

@@ -0,0 +1,51 @@
+
+import time
+import numpy
+from lib.training_data import TrainingDataGenerator, stack_images
+
+
+class Trainer():
+    random_transform_args = {
+        'rotation_range': 10,
+        'zoom_range': 0.05,
+        'shift_range': 0.05,
+        'random_flip': 0.4,
+    }
+
+    def __init__(self, model, fn_A, fn_B, batch_size=64):
+        self.batch_size = batch_size
+        self.model = model
+
+        generator = TrainingDataGenerator(self.random_transform_args, 160)
+        self.images_A = generator.minibatchAB(fn_A, self.batch_size)
+        self.images_B = generator.minibatchAB(fn_B, self.batch_size)
+
+    def train_one_step(self, iter, viewer):
+        epoch, warped_A, target_A = next(self.images_A)
+        epoch, warped_B, target_B = next(self.images_B)
+
+        loss_A = self.model.autoencoder_A.train_on_batch(warped_A, target_A)
+        loss_B = self.model.autoencoder_B.train_on_batch(warped_B, target_B)
+        print("[{0}] [#{1:05d}] loss_A: {2:.5f}, loss_B: {3:.5f}".format(time.strftime("%H:%M:%S"), iter, loss_A, loss_B),
+            end='\r')
+
+        if viewer is not None:
+            viewer(self.show_sample(target_A[0:14], target_B[0:14]), "training")
+
+    def show_sample(self, test_A, test_B):
+        figure_A = numpy.stack([
+            test_A,
+            self.model.autoencoder_A.predict(test_A),
+            self.model.autoencoder_B.predict(test_A),
+        ], axis=1)
+        figure_B = numpy.stack([
+            test_B,
+            self.model.autoencoder_B.predict(test_B),
+            self.model.autoencoder_A.predict(test_B),
+        ], axis=1)
+
+        figure = numpy.concatenate([figure_A, figure_B], axis=0)
+        figure = figure.reshape((4, 7) + figure.shape[1:])
+        figure = stack_images(figure)
+
+        return numpy.clip(figure * 255, 0, 255).astype('uint8')

plugins/Model_IAE/__init__.py

@@ -0,0 +1,8 @@
+# -*- coding: utf-8 -*-
+
+__author__ = """acsaga"""
+__version__ = '0.1.0'
+
+from .Model import Model
+from .Trainer import Trainer
+from .AutoEncoder import AutoEncoder