import os

import gradio as gr
from PIL import Image

from index import query_character


def _fn(image: Image.Image, count: int = 5):
    return query_character(image, count)


if __name__ == '__main__':
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                gr_input_image = gr.Image(type='pil', label='Original Image')
                gr_max_count = gr.Slider(minimum=1, maximum=20, step=1, value=5, label='Max Query Count')
                gr_submit = gr.Button(value='Submit', variant='primary')

            with gr.Column():
                with gr.Tabs():
                    with gr.Tab('Gallery'):
                        gr_gallery = gr.Gallery(label='Gallery')

                    with gr.Tab('Information'):
                        gr_table = gr.DataFrame(label='Information')

            gr_submit.click(
                _fn,
                inputs=[gr_input_image, gr_max_count],
                outputs=[gr_gallery, gr_table],
            )

    demo.queue(os.cpu_count()).launch()