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)