import gradio as gr
from bald_face import BaldFace

bald = BaldFace()


def predict(image):
    return bald.make(image)


footer = r"""
<center>
<b>
Demo for <a href=''>Lightweight OpenPose</a>
</b>
</center>
"""

with gr.Blocks(title="") as app:
    gr.HTML("<center><h1></h1></center>")
    gr.HTML("<center><h3></h3></center>")
    with gr.Row(equal_height=False):
        with gr.Column():
            input_img = gr.Image(type="numpy", label="Input image")
            run_btn = gr.Button(variant="primary")
        with gr.Column():
            output_img = gr.Image(type="pil", label="Output image")
            gr.ClearButton(components=[input_img, output_img], variant="stop")

    run_btn.click(predict, [input_img], [output_img])

    with gr.Row():
        blobs = [[f"examples/{x:02d}.jpg"] for x in range(1, 5)]
        examples = gr.Dataset(components=[input_img], samples=blobs)
        examples.click(lambda x: x[0], [examples], [input_img])

    with gr.Row():
        gr.HTML(footer)

app.launch(share=False, debug=True, show_error=True)
app.queue()