import gradio as gr
import tensorflow as tf
import numpy as np
from matplotlib import cm
from PIL import Image
import imageio

generator = tf.keras.models.load_model('dc_gan.h5')

def interpolate(steps,fps):
    
    #CHANGE LATER
    start = tf.random.normal(shape=(1,128))
    end = tf.random.normal(shape=(1,128))
    #---------------
    
    input_vectors = np.squeeze(np.linspace(start,end,steps))
    
    image_vectors = np.array(generator.predict(input_vectors))
    
    writer = imageio.get_writer('test.mp4', fps=fps)
    
    for im in image_vectors:
        writer.append_data((im*255).astype('uint8'))
    writer.close()
    
    return gr.Video(value = 'test.mp4')


demo = gr.Blocks()
with demo:
    output_interpolation = gr.Video()
    STEPS = gr.Slider(1, 100, step=1,label="Steps")
    FPS = gr.Slider(1, 50, step=1,label="fps")
    btn = gr.Button("Submit")
    btn.click(interpolate, inputs=[STEPS ,FPS], outputs=[output_interpolation])

demo.launch()