import numpy as np
import tensorflow as tf
import gradio as gr
from huggingface_hub import from_pretrained_keras

model = from_pretrained_keras("keras-io/conv_autoencoder")

examples = [
            ['./example_0.jpeg'], 
            ['./example_1.jpeg'], 
            ['./example_2.jpeg'], 
            ['./example_3.jpeg'], 
            ['./example_4.jpeg']
]

def infer(original_image):
    image = tf.keras.utils.img_to_array(original_image)
    image = image.astype("float32") / 255.0
    image = np.reshape(image, (1, 28, 28, 1))
    output = model.predict(image)
    output = np.reshape(output, (28, 28, 1))
    output_image = tf.keras.preprocessing.image.array_to_img(output)
    return output_image

iface = gr.Interface(
    fn = infer,
    title = "Image Denoising using Convolutional AutoEncoders",
    description = "Keras Implementation of a deep convolutional autoencoder for image denoising",
    inputs = gr.inputs.Image(image_mode='L', shape=(28, 28)),
    outputs = gr.outputs.Image(type = 'pil'),
    examples = examples,
    article = "Author: <a href=\"https://huggingface.co/Blazer007\">Vivek Rai</a>. Based on the keras example from <a href=\"https://keras.io/examples/vision/autoencoder/\">Santiago L. Valdarrama</a> \n Model Link: https://huggingface.co/keras-io/conv_autoencoder",
    ).launch(enable_queue=True, debug = True)