from huggingface_hub import from_pretrained_keras
import matplotlib.pyplot as plt
from math import sqrt, ceil
import tensorflow as tf
import gradio as gr
import numpy as np

model1 = tf.keras.models.load_model("mnist.h5", compile=False)
model2 = from_pretrained_keras("keras-io/WGAN-GP")

title = "WGAN-GP"
description = "Image Generation Using WGAN"
article = """
<p style='text-align: center'>
            <a href='https://keras.io/examples/generative/wgan_gp/' target='_blank'>Keras Example given by A_K_Nain</a>
            <br>
            Space by Gitesh Chawda
        </p>
        """

def Predict(model, num_images):
        random_latent_vectors = tf.random.normal(shape=(int(num_images), 128))
        preds = model(random_latent_vectors) 
        num = ceil(sqrt(num_images))
        images = np.zeros((28*num, 28*num), dtype=float)
        n = 0
        for i in range(num):
            for j in range(num):
                if n == num_images:
                    break
                images[i* 28 : (i+1)*28, j*28 : (j+1)*28] = preds[n, :, :, 0]
                n += 1
        return images


def inference(num_images, select: str):
    if select == 'fmnist':
        result = create_digit_samples(model2, num_images)
    else:
        result = create_digit_samples(model1, num_images)
    return result
    
examples = [[5],[8],[2],[3]]
inputs = [gr.inputs.Number(label="number of images"), gr.inputs.Radio(['fmnist', 'mnist'])]
outputs = gr.outputs.Image(label="Output Image")

interface = gr.Interface(
    fn = inference,
    inputs = inputs,
    outputs = outputs,
    examples = examples,
    description = description,
    title = title,
    article = article
    )

interface.launch(share=True)