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


model = from_pretrained_keras("keras-io/lowlight-enhance-mirnet", compile=False)
examples = ['examples/179.png', 'examples/493.png', 'examples/780.png']


def infer(image):
    image = np.expand_dims(image, axis=0)
    output = model.predict(image)
    output_image = output[0] * 255.0
    output_image = output_image.clip(0, 255)
    output_image = output_image.reshape(
        (np.shape(output_image)[0], np.shape(output_image)[1], 3)
    )
    return output_image


iface = gr.Interface(
    fn=infer,
    inputs=[gr.inputs.Image(label="image", type="pil", shape=(128, 128))],
    outputs="image",
    examples=examples).launch(debug=True)