import gradio as gr
from style_transfer import StyleTransfer

style = StyleTransfer()


def predict(content_image, style_image):
    return style.transfer(content_image, style_image)


footer = r"""
<center>
<b>
Demo for <a href='https://www.tensorflow.org/hub/tutorials/tf2_arbitrary_image_stylization'>Style Transfer</a>
</b>
</center>
"""

coffe = r"""
<center>
<a href="https://www.buymeacoffee.com/leonelhs"> <img 
src="https://img.buymeacoffee.com/button-api/?text=Buy me a 
coffee&emoji=&slug=leonelhs&button_colour=FFDD00&font_colour=000000&font_family=Cookie&outline_colour=000000
&coffee_colour=ffffff" /></a>
</center>
"""

with gr.Blocks(title="Style Transfer") as app:
    gr.HTML("<center><h1>Style Transfer</h1></center>")
    gr.HTML("<center><h3>Fast Style Transfer for Arbitrary Styles</h3></center>")
    with gr.Row(equal_height=False):
        with gr.Column():
            with gr.Row(equal_height=True):
                with gr.Column():
                    content_img = gr.Image(type="filepath", label="Content image")
                with gr.Column():
                    style_img = gr.Image(type="filepath", label="Style image")
            run_btn = gr.Button(variant="primary")
        with gr.Column():
            output_img = gr.Image(type="pil", label="Output image")
            gr.ClearButton(components=[content_img, style_img, output_img], variant="stop")

    run_btn.click(predict, [content_img, style_img], [output_img])

    with gr.Row():
        blobs_c = [[f"examples/contents/{x:02d}.jpg"] for x in range(1, 4)]
        examples_c = gr.Dataset(components=[content_img], samples=blobs_c)
        examples_c.click(lambda x: x[0], [examples_c], [content_img])
    with gr.Row():
        blobs_s = [[f"examples/styles/{x:02d}.jpg"] for x in range(1, 12)]
        examples_s = gr.Dataset(components=[style_img], samples=blobs_s)
        examples_s.click(lambda x: x[0], [examples_s], [style_img])

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

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