Upload model weights

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model_data/cvae_trained.ckpt.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,26 @@
+from cvae import get_encoder, get_decoder, CVAE
+import tensorflow as tf
+import gradio as gr
+import numpy as np
+from matplotlib import cm
+from PIL import Image
+
+IMAGE_SIZE = (64, 64)
+model = CVAE(get_encoder(), get_decoder(), latent_dim=512)
+model.load_weights("model_data/cvae_trained.ckpt")
+
+
+def generate_image(mean, variance):
+
+    sample = np.random.normal(mean, variance, size=512)
+    image = tf.reshape(model.decoder(sample[tf.newaxis, :]), IMAGE_SIZE)
+    image = [Image.fromarray(np.uint8(cm.gray(image)*255))]
+
+    return image
+
+
+title = "variational-autoencoder-faces "
+
+gr.Interface(fn=generate_image, outputs=gr.Gallery(), inputs=[gr.inputs.Slider(default=0, label="mean", maximum=10, minimum=-10, step=.1),
+                                                              gr.inputs.Slider(default=1, label="variance", maximum=20, minimum=0, step=.1)],
+             title=title).launch(inline=False)

blocks.py

@@ -0,0 +1,81 @@
+import tensorflow as tf
+from typing import Any, Tuple
+import tensorflow_addons as tfda
+
+
+class ResidualBlock(tf.keras.layers.Layer):
+
+    def __init__(self, filter_num: int, filter_size: int, seed: Any = None, name=None, padding="default",
+                 instance_normalization: bool = False):
+        super(ResidualBlock, self).__init__(name=name)
+        self.filter_num = filter_num
+        self.filter_size = filter_size
+        self.seed = seed
+        self.padding_type = padding
+
+        self.activation_1 = tf.keras.layers.Activation("linear", trainable=False)
+        if padding == "default":
+            self.conv_1 = tf.keras.layers.Conv2D(filters=self.filter_num, kernel_size=self.filter_size,
+                                                 padding="same", trainable=True)
+        elif padding == "reflect":
+            self.pad_1 = ReflectionPadding2D(padding=(1, 1))
+            self.conv_1 = tf.keras.layers.Conv2D(filters=self.filter_num, kernel_size=self.filter_size,
+                                                 padding="valid", trainable=True)
+        else:
+            raise RuntimeError("Non valid padding type.")
+
+        self.activation_2 = tf.keras.layers.Activation("relu")
+
+        if instance_normalization:
+            self.bn_1 = tfda.layers.InstanceNormalization(trainable=True)
+            self.bn_2 = tfda.layers.InstanceNormalization(trainable=True)
+        else:
+            self.bn_1 = tf.keras.layers.BatchNormalization(trainable=True)
+            self.bn_2 = tf.keras.layers.BatchNormalization(trainable=True)
+
+        if padding == "default":
+            self.conv_2 = tf.keras.layers.Conv2D(filters=self.filter_num, kernel_size=self.filter_size,
+                                                 padding="same", trainable=True)
+
+        elif padding == "reflect":
+            self.pad_2 = ReflectionPadding2D(padding=(1, 1))
+            self.conv_2 = tf.keras.layers.Conv2D(filters=self.filter_num, kernel_size=self.filter_size,
+                                                 padding="valid", trainable=True)
+        else:
+            raise RuntimeError("Non valid padding type.")
+
+        self.activation_3 = tf.keras.layers.Activation("relu")
+
+    def call(self, inputs, *args, **kwargs):
+
+        identity = self.activation_1(inputs)
+        x = identity
+        if self.padding_type == "reflect":
+            x = self.pad_1(x)
+        x = self.conv_1(x)
+        x = self.activation_2(x)
+        x = self.bn_1(x)
+        if self.padding_type == "reflect":
+            x = self.pad_2(x)
+        x = self.conv_2(x)
+        x = self.bn_2(x)
+        residual = tf.keras.layers.Add()([x, identity])
+        x = self.activation_3(residual)
+        return x
+
+
+class ReflectionPadding2D(tf.keras.layers.Layer):
+
+    def __init__(self, padding: Tuple[int, int]):
+        super(ReflectionPadding2D, self).__init__()
+        self.pad_width, self.pad_height = padding
+
+    def call(self, inputs, *args, **kwargs):
+        padding_tensor = tf.constant([
+            [0, 0],  # Batch
+            [self.pad_height, self.pad_height],  # Height
+            [self.pad_width, self.pad_width],  # Width
+            [0, 0]  # Channels
+        ])
+
+        return tf.pad(inputs, padding_tensor, mode="REFLECT")

cvae.py

@@ -0,0 +1,127 @@
+import os
+
+import keras.regularizers
+import tensorflow as tf
+from keras.layers import InputLayer, Conv2D, Flatten, BatchNormalization, Dense, UpSampling2D, Reshape, Dropout, Add
+import keras.backend as tfkbk
+import numpy as np
+from blocks import ResidualBlock
+from keras.layers import LeakyReLU, PReLU
+
+INPUT_SHAPE = (64, 64)
+LATENT_DIM = 512
+
+
+def get_encoder():
+    encoder = tf.keras.Sequential(name="encoder")
+
+    encoder.add(InputLayer(input_shape=(*INPUT_SHAPE, 1)))
+
+    encoder.add(Conv2D(32, 3, activation=PReLU(), padding='same', kernel_initializer='he_uniform'))
+    encoder.add(Conv2D(32, 3, activation=PReLU(), padding='same', strides=2, kernel_initializer='he_uniform'))
+    encoder.add(Conv2D(64, 3, activation=PReLU(), padding='same', kernel_initializer='he_uniform'))
+    encoder.add(Conv2D(64, 3, activation=PReLU(), padding='same', strides=2, kernel_initializer='he_uniform'))
+    encoder.add(Conv2D(128, 3, activation=PReLU(), padding='same', kernel_initializer='he_uniform'))
+    encoder.add(Conv2D(128, 3, activation=PReLU(), padding='same', strides=2, kernel_initializer='he_uniform'))
+
+    encoder.add(Flatten())
+
+    encoder.add(Dense(LATENT_DIM * 2, activation=PReLU(), activity_regularizer=tf.keras.regularizers.L2(10e-6)))
+
+    return encoder
+
+
+def get_decoder():
+
+    inputs = tf.keras.layers.Input(shape=[LATENT_DIM, ])
+
+    x = inputs
+    x = Dense(8 * 8 * 16, activation='relu')(x)
+    x = Dense(8 * 8 * 16, activation='relu')(x)
+    x = Reshape(target_shape=(8, 8, 16))(x)
+
+    x = UpSampling2D(2)(x)
+    x = Conv2D(128, 3, activation=LeakyReLU(), padding='same', kernel_initializer='he_uniform')(x)
+    x = ResidualBlock(128, 3, seed=42, name="res1", padding="reflect")(x)
+    x = ResidualBlock(128, 3, seed=42, name="res2", padding="reflect")(x)
+
+    x = UpSampling2D(2)(x)
+    x = Conv2D(64, 3, activation=LeakyReLU(), padding='same', kernel_initializer='he_uniform')(x)
+    x = ResidualBlock(64, 3, seed=42, name="res4", padding="reflect")(x)
+    x = ResidualBlock(64, 3, seed=42, name="res5", padding="reflect")(x)
+
+    x = UpSampling2D(2)(x)
+    x = Conv2D(32, 3, activation=LeakyReLU(), padding='same', kernel_initializer='he_uniform')(x)
+    x = ResidualBlock(32, 3, seed=42, name="res7", padding="reflect")(x)
+    x = ResidualBlock(32, 3, seed=42, name="res8", padding="reflect")(x)
+
+    x = Conv2D(1, 3, padding='same', kernel_initializer='he_uniform')(x)
+
+    return tf.keras.Model(inputs=inputs, outputs=x)
+
+
+class CVAE(tf.keras.Model):
+    def __init__(self, encoder: tf.keras.models.Model, decoder: tf.keras.models.Model,
+                 latent_dim, kl_weight=1, loss_fun='bce', include_regularization: bool = False):
+        super(CVAE, self).__init__()
+        self.kl_weight = kl_weight
+        self.latent_dim = latent_dim
+        self.loss_fun = loss_fun
+        self.encoder = encoder
+        self.decoder = decoder
+        self.kl_loss = 0
+        self.reconstruction_loss = 0
+        self.include_regularization = include_regularization
+
+    def call(self, inputs, training=None, mask=None):
+        z_mean, z_log_var = tf.split(self.encoder(inputs), num_or_size_splits=2, axis=1)
+        z = self.sampling(z_mean, z_log_var, self.latent_dim)
+        # z_mean, z_log_var, z = self.encoder(inputs)
+        outputs = self.decoder(z)
+
+        if training:
+            regularization_loss = tf.math.reduce_sum(self.encoder.losses)
+
+            if self.loss_fun == 'elbo':
+                cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=outputs, labels=inputs)
+                logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])
+                logpz = self.log_normal_pdf(z, 0., 0.)
+                logqz_x = self.log_normal_pdf(z, z_mean, z_log_var)
+                vae_loss = -tf.reduce_mean(logpx_z + logpz - logqz_x)
+            else:
+                kl_loss = 1 + z_log_var - tf.math.square(z_mean) - tf.math.exp(z_log_var)
+                kl_loss = tf.math.reduce_sum(kl_loss, axis=-1)
+                kl_loss *= -0.5 * self.kl_weight
+                self.kl_loss = kl_loss
+                if self.loss_fun == 'mse':
+                    reconstruction_loss = tf.keras.metrics.mean_squared_error(tfkbk.flatten(inputs),
+                                                                              tfkbk.flatten(outputs))
+                elif self.loss_fun == 'bce':
+                    reconstruction_loss = tf.keras.metrics.binary_crossentropy(tfkbk.flatten(inputs),
+                                                                               tfkbk.flatten(outputs))
+                else:
+                    raise ValueError
+
+                reconstruction_loss *= (inputs.shape[1] * inputs.shape[1])
+                self.reconstruction_loss = reconstruction_loss
+                vae_loss = tf.math.reduce_mean(reconstruction_loss + kl_loss)
+
+            if self.include_regularization:
+                vae_loss += regularization_loss
+
+            self.add_loss(vae_loss)
+        return outputs
+
+    @staticmethod
+    def sampling(z_mean, z_log_var, latent_dim):
+        batch = tf.shape(z_mean)[0]
+        epsilon = tf.keras.backend.random_normal(shape=(batch, latent_dim))
+        return z_mean + tf.exp(0.5 * z_log_var) * epsilon
+
+    @staticmethod
+    def log_normal_pdf(sample, mean, logvar, raxis=1):
+        log2pi = tf.math.log(2. * np.pi)
+        return tf.reduce_sum(
+            -.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),
+            axis=raxis)
+

model_data/checkpoint

@@ -0,0 +1,2 @@
+model_checkpoint_path: "cvae_trained.ckpt"
+all_model_checkpoint_paths: "cvae_trained.ckpt"

model_data/cvae_trained.ckpt.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a38425718ad0bb40c171e698c686e9956340f3f3711751df0ab199a36bcdd8a5
+size 137170643

requirements.txt

@@ -0,0 +1,6 @@
+tensorflow~=2.10.0
+tensorflow_addons
+gradio~=3.17.1
+numpy~=1.21.6
+Pillow~=8.4.0
+keras~=2.10.0